'''
* This is the projet for brtc R&D  Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description MCP协议定义模块，包含消息结构、状态码、操作类型等核心定义
* @File: mcp_protocol.py
* @Time: 2025/01/27 14:30:00
* @All Rights Reserve By Brtc
'''

import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional, Callable
from enum import Enum


class OperationType(Enum):
    """协议操作类型枚举"""
    CREATE = "CREATE"
    READ = "READ"
    UPDATE = "UPDATE"
    DELETE = "DELETE"
    SEARCH = "SEARCH"
    COPY = "COPY"
    LIST_TOOLS = "LIST_TOOLS"
    CALL_TOOL = "CALL_TOOL"


class StatusCode(Enum):
    """状态码枚举"""
    SUCCESS = 200
    CREATED = 201
    BAD_REQUEST = 400
    NOT_FOUND = 404
    CONFLICT = 409
    INTERNAL_ERROR = 500


class ContextScope(Enum):
    """上下文作用域枚举"""
    USER = "user"
    SESSION = "session"
    TASK = "task"


class MCPMessage:
    """MCP协议消息类"""
    
    def __init__(
        self,
        operation: OperationType,
        context_id: Optional[str] = None,
        scope: Optional[str] = None,
        payload: Optional[Dict[str, Any]] = None,
        message_id: Optional[str] = None,
        timestamp: Optional[str] = None
    ):
        """
        * @Author Leon-liao
        * @Function: __init__
        * @Description //初始化MCP消息对象
        * @Date :2025/01/27 14:30:00
        * @Param: operation:操作类型枚举, context_id:上下文ID, scope:作用域, 
                 payload:消息负载数据, message_id:消息ID, timestamp:时间戳
        """
        self.header = {
            "protocol": "MCP",
            "version": "1.0",
            "message_id": message_id or f"msg_{uuid.uuid4().hex[:12]}",
            "timestamp": timestamp or datetime.utcnow().isoformat() + "Z"
        }
        self.context = {
            "id": context_id or f"ctx_{uuid.uuid4().hex[:12]}",
            "operation": operation.value,
            "scope": scope or ContextScope.SESSION.value
        }
        self.payload = payload or {}
    
    def to_dict(self) -> Dict[str, Any]:
        """
        * @Author Leon-liao
        * @Function: to_dict
        * @Description //将消息对象转换为字典格式
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 返回字典格式的消息数据
        """
        return {
            "header": self.header,
            "context": self.context,
            "payload": self.payload
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'MCPMessage':
        """
        * @Author Leon-liao
        * @Function: from_dict
        * @Description //从字典创建MCP消息对象
        * @Date :2025/01/27 14:30:00
        * @Param: data:字典格式的消息数据
        * @return: 返回MCP消息对象
        """
        operation = OperationType(data["context"]["operation"])
        return cls(
            operation=operation,
            context_id=data["context"].get("id"),
            scope=data["context"].get("scope"),
            payload=data.get("payload", {}),
            message_id=data["header"].get("message_id"),
            timestamp=data["header"].get("timestamp")
        )


class MCPResponse:
    """MCP协议响应类"""
    
    def __init__(
        self,
        status_code: StatusCode,
        message: str = "",
        data: Optional[Dict[str, Any]] = None,
        request_id: Optional[str] = None
    ):
        """
        * @Author Leon-liao
        * @Function: __init__
        * @Description //初始化MCP响应对象
        * @Date :2025/01/27 14:30:00
        * @Param: status_code:状态码枚举, message:响应消息, data:响应数据, 
                 request_id:请求ID
        * @return: 无
        """
        self.status_code = status_code.value
        self.message = message
        self.data = data or {}
        self.request_id = request_id
        self.timestamp = datetime.utcnow().isoformat() + "Z"
    
    def to_dict(self) -> Dict[str, Any]:
        """
        * @Author Leon-liao
        * @Function: to_dict
        * @Description //将响应对象转换为字典格式
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 返回字典格式的响应数据
        """
        return {
            "status_code": self.status_code,
            "message": self.message,
            "data": self.data,
            "request_id": self.request_id,
            "timestamp": self.timestamp
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'MCPResponse':
        """
        * @Author Leon-liao
        * @Function: from_dict
        * @Description //从字典创建MCP响应对象
        * @Date :2025/01/27 14:30:00
        * @Param: data:字典格式的响应数据
        * @return: 返回MCP响应对象
        """
        return cls(
            status_code=StatusCode(data["status_code"]),
            message=data.get("message", ""),
            data=data.get("data", {}),
            request_id=data.get("request_id")
        )


class Context:
    """上下文数据结构类"""
    
    def __init__(
        self,
        context_id: Optional[str] = None,
        content: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        extensions: Optional[Dict[str, Any]] = None
    ):
        """
        * @Author Leon-liao
        * @Function: __init__
        * @Description //初始化上下文对象
        * @Date :2025/01/27 14:30:00
        * @Param: context_id:上下文ID, content:上下文内容, metadata:元数据, 
                 extensions:扩展字段
        * @return: 无
        """
        self.id = context_id or f"ctx_{uuid.uuid4().hex[:12]}"
        self.metadata = metadata or {}
        self.content = content
        self.extensions = extensions or {}
        
        # 设置默认元数据
        if "created" not in self.metadata:
            self.metadata["created"] = datetime.utcnow().isoformat() + "Z"
        if "modified" not in self.metadata:
            self.metadata["modified"] = datetime.utcnow().isoformat() + "Z"
    
    def update_modified_time(self):
        """
        * @Author Leon-liao
        * @Function: update_modified_time
        * @Description //更新最后修改时间
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 无
        """
        self.metadata["modified"] = datetime.utcnow().isoformat() + "Z"
    
    def to_dict(self) -> Dict[str, Any]:
        """
        * @Author Leon-liao
        * @Function: to_dict
        * @Description //将上下文对象转换为字典格式
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 返回字典格式的上下文数据
        """
        return {
            "id": self.id,
            "metadata": self.metadata,
            "content": self.content,
            "extensions": self.extensions
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Context':
        """
        * @Author Leon-liao
        * @Function: from_dict
        * @Description //从字典创建上下文对象
        * @Date :2025/01/27 14:30:00
        * @Param: data:字典格式的上下文数据
        * @return: 返回上下文对象
        """
        return cls(
            context_id=data.get("id"),
            content=data.get("content"),
            metadata=data.get("metadata", {}),
            extensions=data.get("extensions", {})
        )


class ToolParameter:
    """工具参数定义类"""
    
    def __init__(
        self,
        name: str,
        param_type: str,
        description: str,
        required: bool = True,
        default: Optional[Any] = None
    ):
        """
        * @Author Leon-liao
        * @Function: __init__
        * @Description //初始化工具参数对象
        * @Date :2025/01/27 14:30:00
        * @Param: name:参数名称, param_type:参数类型, description:参数描述, 
                 required:是否必需, default:默认值
        * @return: 无
        """
        self.name = name
        self.type = param_type
        self.description = description
        self.required = required
        self.default = default
    
    def to_dict(self) -> Dict[str, Any]:
        """
        * @Author Leon-liao
        * @Function: to_dict
        * @Description //将参数对象转换为字典格式
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 返回字典格式的参数数据
        """
        result = {
            "name": self.name,
            "type": self.type,
            "description": self.description,
            "required": self.required
        }
        if self.default is not None:
            result["default"] = self.default
        return result
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ToolParameter':
        """
        * @Author Leon-liao
        * @Function: from_dict
        * @Description //从字典创建工具参数对象
        * @Date :2025/01/27 14:30:00
        * @Param: data:字典格式的参数数据
        * @return: 返回工具参数对象
        """
        return cls(
            name=data["name"],
            param_type=data["type"],
            description=data["description"],
            required=data.get("required", True),
            default=data.get("default")
        )


class Tool:
    """工具定义类"""
    
    def __init__(
        self,
        name: str,
        description: str,
        parameters: Optional[List[ToolParameter]] = None,
        handler: Optional[Callable[..., Any]] = None
    ):
        """
        * @Author Leon-liao
        * @Function: __init__
        * @Description //初始化工具对象
        * @Date :2025/01/27 14:30:00
        * @Param: name:工具名称, description:工具描述, parameters:参数列表, 
                 handler:工具处理函数
        * @return: 无
        """
        self.name = name
        self.description = description
        self.parameters = parameters or []
        self.handler = handler
    
    def to_dict(self) -> Dict[str, Any]:
        """
        * @Author Leon-liao
        * @Function: to_dict
        * @Description //将工具对象转换为字典格式
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 返回字典格式的工具数据
        """
        return {
            "name": self.name,
            "description": self.description,
            "parameters": [param.to_dict() for param in self.parameters]
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Tool':
        """
        * @Author Leon-liao
        * @Function: from_dict
        * @Description //从字典创建工具对象
        * @Date :2025/01/27 14:30:00
        * @Param: data:字典格式的工具数据
        * @return: 返回工具对象
        """
        parameters = [
            ToolParameter.from_dict(p) for p in data.get("parameters", [])
        ]
        return cls(
            name=data["name"],
            description=data["description"],
            parameters=parameters
        )

